#!/bin/bash
#
# Copyright:: Copyright (c) 2015 GitLab B.V.
# License:: Apache License, Version 2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

# This wrapper is meant to be invoked by omnibus-gitlab via Runit

# Let runit capture all script error messages
exec 2>&1

function assert_non_empty
{
  if [ -z "$(eval echo \$$1)" ] ; then
    echo "$0 error: \$$1 is empty"
    exit 1
  fi
}

# We expect the following variables to be set for us in the environment
assert_non_empty current_pidfile
assert_non_empty rails_app
assert_non_empty user
assert_non_empty environment
assert_non_empty unicorn_rb

readonly oldbin_pidfile=${current_pidfile}.oldbin
readonly unicorn_wait_start=1 # time in seconds
readonly unicorn_poll_alive=1 # time in seconds
readonly unicorn_kill_timeout=5 # time on seconds
readonly unicorn_wait_reexec=120 # time on seconds
readonly unicorn_wait_old_pid=90 # time on seconds

function main
{
  cd /opt/gitlab/embedded/service/${rails_app}
  find_us_a_unicorn
  trap_signals
  wait_for_unicorn_to_exit
}

function find_us_a_unicorn
{
  adopt ${current_pidfile}
  if [[ ${unicorn_pid} ]]; then
    echo "adopted existing unicorn master ${unicorn_pid}"
    return
  fi

  adopt ${oldbin_pidfile}
  if [[ ${unicorn_pid} ]]; then
    echo "adopted existing oldbin unicorn master ${unicorn_pid}"
    return
  fi

  echo "starting new unicorn master"
  start_unicorn_master
  sleep ${unicorn_wait_start}

  adopt ${current_pidfile}
  if [[ ${unicorn_pid} ]]; then
    echo "adopted new unicorn master ${unicorn_pid}"
    return
  fi

  echo "failed to start a new unicorn master"
  exit
}

function adopt
{
  local pid=$(cat $1 2>/dev/null)
  if alive ${pid} && is_unicorn ${pid}; then
    readonly unicorn_pid=${pid}
  fi
}

function alive
{
  kill -0 $1 > /dev/null 2>&1
}

function is_unicorn
{
  ps -p $1 -o args | grep -q unicorn
}

function start_unicorn_master
{
  chpst -e /opt/gitlab/etc/${rails_app}/env -U ${user}:${group} \
    /opt/gitlab/embedded/bin/bundle exec unicorn \
      -D \
      -E ${environment} \
      -c ${unicorn_rb} \
      /opt/gitlab/embedded/service/${rails_app}/config.ru
}

function trap_signals
{
  # Forward all common runit signals except:
  # - HUP which we handle below;
  # - KILL which cannot be caught.
  for sig in STOP CONT ALRM QUIT USR1 USR2; do
    trap "forward_signal ${sig}" ${sig}
  done

  for sig in INT TERM; do
    trap "forward_signal_and_kill_unicorn ${sig}" ${sig}
  done

  # Omnibus-ctl does not have a subcommand that sends USR2 but it can send HUP.
  # To allow for reloading unicorn from the command line, translate HUP to
  # USR2, then QUIT the old unicorn.
  trap "echo 'wrapper received HUP'; forward_usr2_and_quit_old_unicorn" HUP
}

function forward_signal
{
  echo "forwarding $1 to unicorn master ${unicorn_pid}"
  kill -$1 ${unicorn_pid}
}

function forward_signal_and_kill_unicorn
{
  forward_signal $1
  kill_unicorn_if_running $1
}

function forward_usr2_and_quit_old_unicorn
{
  forward_signal USR2
  quit_old_unicorn_if_running
}

function wait_for_unicorn_to_exit
{
  while sleep ${unicorn_poll_alive}; do
    alive ${unicorn_pid} || break
  done
}

function kill_unicorn_if_running
{
  counter=0
  while (alive ${unicorn_pid}) && [[ $counter -lt ${unicorn_kill_timeout} ]]; do
    echo "Waiting for unicorn to die..."
    let counter++
    sleep ${unicorn_poll_alive}
  done

  # Forcibly kill if still running
  if alive ${unicorn_pid}; then
     echo "Forcibly killing unicorn"
     kill -9 ${unicorn_pid}
  fi

  # If a new process has been reexecuted, forward the same signal
  # and forcibly kill it if it does not respond.
  if [ ! -f ${current_pidfile} ]; then
    return 0
  fi

  local new_unicorn_pid=$(cat $current_pidfile 2>/dev/null)

  if alive ${new_unicorn_pid}; then
    echo "Sending $1 signal to new unicorn..."
    kill -$1 ${new_unicorn_pid}
  fi

  counter=0

  while (alive ${new_unicorn_pid}) && [[ $counter -lt ${unicorn_wait_reexec} ]]; do
    echo "Waiting for new unicorn to die..."
    let counter++
    sleep ${unicorn_poll_alive}
  done

  # Forcibly kill if still running
  if alive ${new_unicorn_pid}; then
     echo "Forcibly killing new unicorn"
     kill -9 ${new_unicorn_pid}
  fi
}

function quit_old_unicorn_if_running
{
  local old_unicorn_pid=$(cat $oldbin_pidfile 2>/dev/null)
  counter=0

  while [ ! -f ${oldbin_pidfile} ]; do
    echo "Waiting for old unicorn PID to appear..."

    sleep 1
    old_unicorn_pid=$(cat $oldbin_pidfile 2>/dev/null)
    let counter++

    if [[ $counter -gt ${unicorn_wait_old_pid} ]]; then
      break
    fi
  done

  if [ ! -f ${oldbin_pidfile} ]; then
    echo "No old PID found..."
    return 1
  fi

  counter=0

  while [ ! -f ${current_pidfile} ]; do
    echo "Waiting for new unicorn to reexec..."

    sleep 1
    let counter++

    if [[ $counter -gt ${unicorn_wait_reexec} ]]; then
      break
    fi
  done

  counter=0
  local new_pid=$(cat $current_pidfile 2>/dev/null)

  while [[ $(ps --no-headers -o pid --ppid=$new_pid | wc -w) -lt 1 ]]; do
    echo "Waiting until $new_pid forks new workers..."

    sleep 1
    let counter++

    if [[ $counter -gt ${unicorn_wait_reexec} ]]; then
      break
    fi
  done

  sleep 1

  # At this point, the new unicorn should be alive. It could be that the
  # old unicorn died anyway, so let's not keep waiting.
  if alive ${old_unicorn_pid}; then
    echo "Sending QUIT signal to old unicorn..."
    kill -QUIT ${old_unicorn_pid}
  fi

  counter=0

  while (alive ${old_unicorn_pid}) && [[ $counter -lt ${unicorn_wait_reexec} ]]; do
    echo "Waiting for old unicorn to die..."
    let counter++
    sleep ${unicorn_poll_alive}
  done

  # Forcibly kill if still running
  if alive ${old_unicorn_pid}; then
     echo "Forcibly killing old unicorn"
     kill -9 ${old_unicorn_pid}
  fi
}

main
echo "wrapper for unicorn master ${unicorn_pid} exiting"
